# this is for primer design for tiling long regions for MIPs

import os
import sys
sys.path.append(os.path.join(os.getcwd(), "modules"))
import cPickle
import random
import pysam

from itertools import product
# now import custom functions
from allfunctions import revcomp, Tm, dimer, GC
from modefunctions import find_exon, check_snp

def pck_small(outputdir, filename, pnumber, header, seq, runDB, pathDB, Random):
    samp = 200

    designmode = runDB["designmode"]
    POI = runDB["POI"]
    
    ############## DEFINE FUNCTIONS ############## 

    # this is the picker function 
    # the function input is a list
    # this is the picker function 
    # the function input is a list
    if designmode == 'PCR':
        from modefunctions import pcr_picker as picker
        #dna=250.0; Na=50.0; K=0; Tris=0; Mg=1.5; dNTP=0.2; DMSO=0

        Arm_max = "na" # total length of forward and reverse arms allowed
        REdict = "na"
        PriF = "na" # Forward MIPamp primer site
        Middle = "na" # The backbone
        PriR = "na" # Reverse MIPamp primer site
        Oligo_len = "na" # Lenght of the final product for ordering

        Remove = False
        if runDB["repeats"] == "ignore":
            from modefunctions import pcr_score_pair_IR as score_pair
        elif runDB["repeats"] == "score" :
            from modefunctions import pcr_score_pair_S as score_pair
        elif runDB["repeats"] == "remove":
            from modefunctions import pcr_score_pair_IR as score_pair
            Remove = True

    elif designmode == 'MIP':
        from modefunctions import mip_picker as picker
        #dna=10; Na=5; K=27; Tris=22.9; Mg=10; dNTP=0.3; DMSO=0
        # use values from ampligase buffer, D1/20 Ampligase storage buffer and, D1/50 Phusion storage buffer

        Arm_max = runDB['Arm_max'] # total length of forward and reverse arms allowed
        REdict= runDB['REdict'] 
        PriF = runDB['PriF'].lower() # Forward MIPamp primer site
        Middle = runDB['Middle'].lower() # The backbone
        PriR = runDB['PriR'].lower() # Reverse MIPamp primer site
        Oligo_len = runDB['Oligo_len'] # Lenght of the final product for ordering

        Remove = False
        if runDB["repeats"] == "ignore":
            from modefunctions import score_mip_IR as score_pair
        elif runDB["repeats"] == "score":
            from modefunctions import score_mip_S as score_pair
        elif runDB["repeats"] == "remove":
            from modefunctions import score_mip_IR as score_pair
            Remove = True

    ### END OF FUNCTIONS ####


    ##########################################
    # ENTER RUN SPECIFIC CRITERIA
    MinGC, MaxGC = runDB['MinGC'], runDB['MaxGC']
    MinTm, MaxTm = runDB['MinTm'], runDB['MaxTm']
    MinLen, MaxLen = runDB['MinLen'], runDB['MaxLen']
    Epos = runDB['Epos'] # relative position of the exon in the target (each end)
    Smin, Smax = runDB['Smin'], runDB['Smax'] 
    Olap = runDB['Olap'] # tiling overlap
    pntcutoff = runDB['pntcutoff']
    dna = runDB["dna"]
    Na = runDB["Na"]
    K = runDB["K"]
    Tris = runDB["Tris"]
    Mg = runDB["Mg"]
    dNTP = runDB["dNTP"]
    DMSO = runDB["DMSO"]
    tiling = runDB["tiling"]
    #########################################

    # all the output file data will be collected in a dict to be PIPE_OUT
    output_data = {"display":{pnumber: list()}, "picked":list()}

    # start the processing the header and the seq
    # pnumber for primer naming (K)
    # Pi is to mark each pair with a unique name
    Pi = 1

    failed = "NA" # count number of primer failed by SNP

    bases = set(['a', 'c', 'g', 't'])
    k = 1 # counter for primer F
    t = 1 # counter for primer R
    # l for primer start location
    # m for min primer length
    # M for primer stop location
    # n  for max primer length

    R = list() # reverse primers list
    F = list() # forwrd primers list
    Picked_code = dict()

    genome, Chr, Tstart, Tstop, Offset = header

    repeatDB = pathDB["repeatDB"]
    exonfile = pysam.Tabixfile(os.path.join(pathDB["exonDB"], "%s.tabix.gz" %Chr))
    snpfile = pysam.Tabixfile(os.path.join(pathDB["snpDB"], "%s.tabix.gz" %Chr))

    Epos = runDB["Epos"] + Offset
    Estart = Tstart + Epos
    Estop = Tstop - Epos
    target = '%s:%s-%s' %(Chr, Estart, Estop)
    len_seq = len(seq)

    scale = max((Tstop - Tstart) / 100, 1)

    # to be used in the display file for scaling
    output_data["display"][pnumber].append('\n%s\t%s:%s-%s\n' %tuple(header[0:4]))
    output_data["display"][pnumber].append('\nRegion %s Scale is %d\n' %(pnumber, scale))
    
    if not POI:
        output_data["display"][pnumber].append('%s%s%s\n' %('-'*(Epos/scale), '|'*((Estop-Estart)/scale), '-'*(Epos/scale)))
        # define seqR
        seqR = revcomp(seq)
    
    else:
        output_data["display"][pnumber].append('%s*%s\n' %('|'*(len_seq/(2*scale)), '|'*(len_seq/(2*scale))))
        # if poi then F has to be on the left, R has to be on the right
        seqR = revcomp(seq[Epos + 2 * MaxLen : ])
        seq = seq[ : Epos - 2 * MaxLen + 1]
        # so do not waste time designing others

    l = -1
    m = MinLen-2
    n = MaxLen
    site = len(seq)
    while l < site-(n+1):
        l += 1
        i = 0
        m += 1
        M = m+1
        while i < (n - MinLen + 1): 
            Fprimer = seq[l:M]
            len_Fprimer = len(Fprimer)
            Fstart = Tstart+l
            Fstop = Tstart+M-1
            if Fstop <= Tstop-Smin and sum([Fprimer.count(p) for p in bases]) == len_Fprimer:
                # the primer should not start near the Tstop, unrelevant
                # the primer should not have any other base then a,c,g,t
                tm = Tm(Fprimer, dna, Na, K, Tris, Mg, dNTP, DMSO)
                gc = GC(Fprimer)
                if MinTm <= tm <= MaxTm and MinGC <= gc <= MaxGC:
                    if check_snp(snpfile, Chr, Fstart, Fstop, "F") == 'PASS':
                        if dimer(Fprimer, Fprimer, pntcutoff=pntcutoff) == 'PASS':
                            name = '%sF%s' %(pnumber, k)
                            F.append([name, Fprimer, tm, gc, len_Fprimer, target, Chr, Fstart, Fstop, Epos-M])
                            k += 1
            Rprimer = seqR[l:M]
            len_Rprimer = len(Rprimer)
            Rstart = Tstop-M+1
            Rstop = Tstop-l
            if Rstart >= Tstart+Smin and sum([Rprimer.count(p) for p in bases]) == len_Rprimer:
                # the primer should not start near the Tstop, unrelevant
                # the primer should not have any other base then a,c,g,t
                tm = Tm(Rprimer, dna, Na, K, Tris, Mg, dNTP, DMSO)
                gc = GC(Rprimer)
                if MinTm <= tm <= MaxTm and MinGC <= gc <= MaxGC:
                    if check_snp(snpfile, Chr, Rstart, Rstop, "R") == 'PASS':
                        if dimer(Rprimer, Rprimer, pntcutoff=pntcutoff) == 'PASS':
                            name= '%sR%s' %(pnumber, t)
                            R.append([name, Rprimer, tm, gc, len_Rprimer, target, Chr, Rstart, Rstop, Epos-M])
                            t+=1
            M+=1
            i+=1

    # THE TILER FUNCTION
    # here goes the picking function for tiling
    RstartN = Tstart
    FstopN = Tstart
    Lsize = Smin
    TTemp2 = list()
    len_TTemp2 = 0
    pick_pass = 0
    last_run = False
    pick_reverse = False
    while 1: #there is a break statement when conditions are met
        if Random:
            if designmode == 'PCR':
                TTemp= [[0] + pair[0] + pair[1] + [(pair[1][8]-pair[0][7]+1), '%sP%s' %(pnumber, Pi)] for Pi, pair in enumerate(product(random.sample(F, min(len(F), samp)), random.sample(R, min(len(R), samp))), Pi) if Smin <= (pair[1][8]-pair[0][7]+1) <= Smax]
            elif designmode == 'MIP':
                TTemp= [[0] + pair[0] + pair[1] + [(pair[1][8]-pair[0][7]+1), '%sP%s' %(pnumber, Pi)] for Pi, pair in enumerate(product(random.sample(F, min(len(F), samp)), random.sample(R, min(len(R), samp))), Pi) if Smin <= (pair[1][8]-pair[0][7]+1) <= Smax and pair[0][4] + pair[1][4] <= Arm_max]

        else:
            # if you are not picking random and
            # you have already defined TTemp
            # no need to redefine
            try:
                TTemp
            except NameError:
                if designmode == 'PCR':
                    TTemp= [[0] + pair[0] + pair[1] + [(pair[1][8]-pair[0][7]+1), '%sP%s' %(pnumber, Pi)] for Pi, pair in enumerate(product(F, R), Pi) if Smin <= (pair[1][8]-pair[0][7]+1) <= Smax]

                elif designmode == 'MIP':
                    TTemp= [[0] + pair[0] + pair[1] + [(pair[1][8]-pair[0][7]+1), '%sP%s' %(pnumber, Pi)] for Pi, pair in enumerate(product(F, R), Pi) if Smin <= (pair[1][8]-pair[0][7]+1) <= Smax and (pair[0][4] + pair[1][4]) <= Arm_max]


        # now start picking
        picking = 'FAIL'
        old_pick = list()

        len_TTemp2= len(TTemp2)
        if len_TTemp2 == 0 or len_TTemp2 % 2 == 0:
            # even
            even= True
        elif len_TTemp2 % 2 != 0:
            # odd
            even= False

        # ARE YOU TILING BASES
        # OR TILING REGIONS
        if runDB["tiling"]:
            """
            if Estop >= RstartN > Tstart:
                TTemp1= [[score_pair(Li, repeatDB, Smax, even)] + Li[1:] for Li in TTemp if FstopN < Li[8] < RstartN - Lsize*(1-1.0/tiling) and Li[18] > RstartN + MaxLen]
                # have a fail safe
                if not len(TTemp1):
                    TTemp1= [[score_pair(Li, repeatDB, Smax, even)] + Li[1:] for Li in TTemp if Li[8] < RstartN - Lsize*(1-1.0/tiling) and Li[18] > RstartN + MaxLen]

            elif RstartN > Estop:
                TTemp1= [[score_pair(Li, repeatDB, Smax, even)] + Li[1:] for Li in TTemp if FstopN < Li[8] < RstartN - Lsize*(1-1.0/tiling) and Li[18] > Estop]
                # there may be cases where a short exon is not tiled
                if not len(TTemp1) and len(TTemp2) == 1:
                    TTemp1= [[score_pair(Li, repeatDB, Smax, even)] + Li[1:] for Li in TTemp if Li[9] < FstopN - MaxLen and Li[18] < RstartN]

                else:
                    # is there something you can do
                    pass

                last_run= True

            else:
                #print "this is %s pick" %(len_TTemp2 + 1)
                # this ensure that first tiling primer is picked
                # within the 5' buffer region outside the target exon
                # always try to cover the region with one product
                if Tstop - Tstart <= Smax:
                    TTemp1= [[score_pair(Li, repeatDB, Smax, even)] + Li[1:] for Li in TTemp if Li[9] < Estart and Li[18] > Estop]
                else:
                    TTemp1= [[score_pair(Li, repeatDB, Smax, even)] + Li[1:] for Li in TTemp if Li[9] < Estart and Li[18] > Estart]
                    """

            if Estop >= RstartN > Tstart:
                TTemp1= [[score_pair(Li, repeatDB, Smax, even)] + Li[1:] for Li in TTemp if FstopN < Li[8] < RstartN - Lsize*(1-1.0/tiling) and Li[18] > RstartN + MaxLen]
                # have a fail safe
                if not len(TTemp1):
                    TTemp1 = [[score_pair(Li, repeatDB, Smax, even)] + Li[1:] for Li in TTemp if Li[8] < RstartN - Lsize*(1-1.0/tiling) and Li[18] > RstartN + MaxLen]
                    

            elif RstartN > Estop:
                TTemp1= [[score_pair(Li, repeatDB, Smax, even)] + Li[1:] for Li in TTemp if FstopN < Li[8] < RstartN - Lsize*(1-1.0/tiling) and Li[18] > Estop]
                if not len(TTemp1):
                    if len (TTemp2) == 1:
                        # make sure you cover all four possible primer picks
                        TTemp1 = []
                        for Li in TTemp:
                            if FstopN < Li[8] < Li[9] < Estart < Estop < Li[18] < Li[19] < RstartN:
                                TTemp1.append([score_pair(Li, repeatDB, Smax, even)] + Li[1:])
                            elif Li[9] < FstopN - MaxLen < Estart < Estop < RstartN + MaxLen < Li[18]:
                                TTemp1.append([score_pair(Li, repeatDB, Smax, even)] + Li[1:])
                            elif Li[9] < FstopN - MaxLen < Estart < Estop < Li[18] < Li[19] < RstartN:
                                TTemp1.append([score_pair(Li, repeatDB, Smax, even)] + Li[1:])
                            elif FstopN < Li[8] < Li[9] < Estart < Estop < RstartN + MaxLen < Li[18]:
                                TTemp1.append([score_pair(Li, repeatDB, Smax, even)] + Li[1:])
                    """
                    else:
                        TTemp1= []
                        for Li in TTemp:
                            if FstopN < Li[8] < Li[19] < RstartN:
                                TTemp1.append([score_pair(Li, repeatDB, Smax, even)] + Li[1:])
                            elif Li[9] < FstopN - MaxLen < RstartN + MaxLen < Li[18]:
                                TTemp1.append([score_pair(Li, repeatDB, Smax, even)] + Li[1:])
                            elif Li[9] < FstopN - MaxLen < Li[18] < Li[19] < RstartN:
                                TTemp1.append([score_pair(Li, repeatDB, Smax, even)] + Li[1:])
                            elif FstopN < Li[8] < Li[9] < RstartN + MaxLen < Li[18]:
                                TTemp1.append([score_pair(Li, repeatDB, Smax, even)] + Li[1:])
                                """

                        
                    # if all else fails, just give me whatever you have
                    if not len(TTemp1):
                        TTemp1= [[score_pair(Li, repeatDB, Smax, even)] + Li[1:] for Li in TTemp if Li[8] < RstartN and Li[18] > Estop]
                last_run= True

            else:
                #print "this is %s pick" %(len_TTemp2 + 1)
                # this ensure that first tiling primer is picked
                # within the 5' buffer region outside the target exon
                # always try to cover the region with one product
                if Tstop - Tstart <= Smax:
                    TTemp1= [[score_pair(Li, repeatDB, Smax, even)] + Li[1:] for Li in TTemp if Li[9] < Estart and Li[18] > Estop]
                else:
                    TTemp1= [[score_pair(Li, repeatDB, Smax, even)] + Li[1:] for Li in TTemp if Li[9] < Estart - MaxLen and Li[18] > Estart]

            #now pick
            #print RstartN, Tstop
            #print len(TTemp1)
            if len(TTemp1) > 0:
                Picked_code, thepair= picker(TTemp1, REdict, PriF, Middle, PriR, Arm_max, pntcutoff, pick_reverse, TTemp2, even, repeatDB, Picked_code, runDB['gfserver'], Remove)

                if thepair != 'FAIL':
                    #print thepair[-2], len(thepair[-2])
                    thepair.extend([find_exon(exonfile, Chr, thepair[8], thepair[19])])

                    TTemp2.append(thepair)
                    # establish the last lenght of TTemp2
                    len_TTemp2= len(TTemp2)
                    picking= 'PASS'

                    # adjust RstartN
                    RstartN= TTemp2[-1][18]
                    Lsize= TTemp2[-1][21]
                    FstopN= TTemp2[-1][9]
                    #this is the Reverse start position of the last pair
                    #this is the Reverse stop position of the last pair 
                    #print thepair

            # modify seach criteria and try again
            if last_run:
                break

            # now its time to go back to the top to move the region
            # in order to determine where the new region will start
            if picking == 'FAIL':
                RstartN += 5
                FstopN += 5

        else:       
            if RstartN > Tstart:
                # any other case
                # although first try to find the longest amplicon that can go all the way
                if Tstop - RstartN + Olap <= Smin:
                    TTemp1= [[score_pair(Li, repeatDB, Smax, even)] + Li[1:] for Li in TTemp if Li[18] > Estop]
                    if not len(TTemp1):
                        # you are miserable, just give me something, work with me
                        TTemp1= [[score_pair(Li, repeatDB, Smax, even)] + Li[1:] for Li in TTemp if Li[8] > RstartN - Smin + 2*MaxLen]
                    pick_reverse= None
                else:
                    TTemp1= [[score_pair(Li, repeatDB, Smax, even)] + Li[1:] for Li in TTemp if RstartN - Olap <= Li[9] < RstartN and Li[18] > RstartN + min(Olap+2*MaxLen, Smax)]

                """
                rule 1:
                the new pairs Fstop cannot be smaller than the Rstart-Olap of the old pick
                but has to be smaller than Rstart of that primer

                rule 2:
                the new pairs Rstart has to be futher away from the Rstart of the old pick
                by Olap+2*MaxLen 
                (one for the offset from Rstop-Rstart, other from Li[9]- Li[8]), 
                so that the new Fstart won't overlap with the previous pairs
                Rstop
                """

            elif POI:
                # if it is a point-of-interest pick that is all you need
                # one good pick
                TTemp1= [[score_pair(Li, repeatDB, Smax, even)] + Li[1:] for Li in TTemp]

            else:
                # this ensure that first tiling primer is picked
                # within the 5' buffer region outside the target exon
                # always try to cover the region with one product
                if Tstop - Tstart <= Smax:
                    TTemp1= [[score_pair(Li, repeatDB, Smax, even)] + Li[1:] for Li in TTemp if Li[9] < Estart and Li[18] > Estop]
                else:
                    TTemp1= [[score_pair(Li, repeatDB, Smax, even)] + Li[1:] for Li in TTemp if Li[9] < Estart and Li[18] > Estart]


            #TTemp1.sort(reverse= True)
            #print TTemp1[0][5] + TTemp1[0][15]

            #now pick
            #print RstartN, Tstop
            #print len(TTemp1)
            if len(TTemp1) > 0:
                Picked_code, thepair= picker(TTemp1, REdict, PriF, Middle, PriR, Arm_max, pntcutoff, pick_reverse, TTemp2, even, repeatDB, Picked_code, runDB['gfserver'], Remove)

                if thepair != 'FAIL':
                    #print thepair[-2], len(thepair[-2])
                    thepair.extend([find_exon(exonfile, Chr, thepair[8], thepair[19])])

                    TTemp2.append(thepair)
                    # establish the last lenght of TTemp2
                    len_TTemp2= len(TTemp2)
                    picking= 'PASS'
                    pick_pass += 1
                    if pick_pass == 2:
                        # you need two consequtive picks before clearing
                        old_pick= []
                        pick_pass= 0
                    # adjust RstartN
                    RstartN= TTemp2[-1][18]
                    #this is the Reverse start position of the last pair
                    #this is the Reverse stop position of the last pair 
                    #print thepair
                    if RstartN > Estop:
                        break

            # else:
            # if len(TTemp1) == 0: it means that there the picking is a FAIL
            # modify seach criteria and try again
            if POI or pick_reverse == None:
                # for POI
                # you get one shot whether it is pass or fail does not matter
                # there is nothing you can do to make it pass after fail
                break

        # now its time to go back to the top to move the region
        # in order to determine where the new region will start
        if picking == 'FAIL':
            if pick_pass == 1 and len(old_pick) == 1:
                # you failed before and picked a new one but you cannot fool me
                # you failed yet again
                # don't forget to remove the new pick
                # but which one to keep, they both are bad
                try:
                    TTemp2.pop()
                except IndexError:
                    # if TTemp2 is empty now
                    pass
                TTemp2.append(old_pick.pop())
                # you tried everything you can but it does not work
                # fix len(TTemp2) so there won't be any endless loops going back
                len_TTemp2= len(TTemp2)
                # the old pick now determines the RstartN
                RstartN= TTemp2[-1][18]
                FstopN= TTemp2[-1][9]
                if Tstop - RstartN + runDB["Olap"] < Smin:
                    # pick one from reverse direction
                    pick_reverse= True
                else:
                    # just a difficult region, move along
                    RstartN+= 5
                    FstopN+= 5


            else:
                # you failed straight up
                if len(TTemp2) > len_TTemp2:
                    # no worries, do over, try a new pick
                    if len(old_pick) == 0:
                        old_pick.append(TTemp2.pop())
                        try:
                            RstartN= TTemp2[-1][18]
                            FstopN= TTemp2[-1][9]
                        except IndexError:
                            RstartN= Tstart
                            FstopN= Tstart

                    else:
                        # you failed even after you tried a new pick
                        TTemp2.append(old_pick.pop())
                        # the old pick now determines the RstartN
                        RstartN= TTemp2[-1][18]
                        FstopN= TTemp2[-1][9]
                        if Tstop - RstartN + runDB["Olap"] < Smin:
                            # pick one from reverse direction
                            pick_reverse= True
                        else: 
                            # it could not inititate primer picking
                            RstartN+= 5
                            FstopN+= 5

                            # use plus one in case minlen= maxlen
                        """
                        IMPORTANT
                        the two if and elif statements may counter each other
                        so make sure that the net effect is going forward
                        so that they don't end up in a loop
                        thus go back MinLen, go forward MaxLen
                        """
                else:
                    # cannot initiate primer picking so move forward
                    if Tstop - RstartN + runDB["Olap"] < Smin:
                        # pick one from reverse direction
                        pick_reverse= True
                    else: 
                        # it could not inititate primer picking
                        RstartN+= 5
                        FstopN+= 5
                                        
                        # use plus one in case minlen= maxlen

        if RstartN > Tstop or RstartN < Tstart:
            # for the rest if you are out-of-bounce you are out
            break

        # generate the outputs
        # reset or adjust counters
        #print 'Process %s is %s complete' % (str(core), str(int(100*ii/tcount))+'%')
        Li= 0
        k, t, = 1, 1
        #END OF THE TILER FUNCTION

        if pick_reverse == None and len(TTemp2) > 1:
            if len(TTemp2) > 2:
                rstartn= TTemp2[-3][18]
            else:
                rstartn= Estart

            # make sure that the last two primers are not reduntant.
            # pick the one that cover the most seq
            l1= TTemp2.pop() # the reverse pick
            l2= TTemp2.pop() # last forward pick
            keep= False
            for each in [l1, l2]:
                if each[9] < rstartn < Estop < each[18]:
                    # who ever that is add it back
                    TTemp2.append(each)
                    keep= True
                    break
            # the previous elimination did not work, try nearest neighbour
            if not keep:
                if abs(l1[8]-l2[8]) <=10 or abs(l1[19]-l2[19]) <= 10:
                    # do not use valuble oligo space for the gain of a few extra bases
                    # pick the one that has the most coverage
                    cdb={}
                    for each in [l1,l2]:
                        covg= each[18]-each[9]
                        # this is unique covg
                        # correct for non-target covg
                        if each[10] > 0:
                            covg-= each[10]
                        if each[20] > 0:
                            covg-= each[20]
                        cdb[covg]= each
                    TTemp2.append(cdb[max(cdb.keys())])
                else:
                    # if nothing else but both back
                    TTemp2.extend([l2, l1])

    # this is common for both tiling bases and tiling regions
    for line in TTemp2:
        output_data["display"][pnumber].append('%s%s%s\n' %(' '*((line[8]-Tstart)/scale), '-'*((line[19]-line[8])/scale), ' '*((Tstop-line[19])/scale)))
        output_data["picked"].append(line)
    output_data["display"][pnumber].append("\n")

    # write out the data file

    """
    with open(os.path.join(outputdir, '%s.pkl' %filename), "w") as output:
        cPickle.dump(output_data, output)
        """
    
    return output_data

